'''
draw pictures
including accuracy, loss, TI
'''

import os
import numpy as np
import tool as tl
import matplotlib.pyplot as plt
from pic_draw.acc import acc_picnum as accpic
from pic_draw.loss import loss_picnum as losspic
from pic_draw.distance import dis_picnum as dispic
import pic_draw.ie_draw as ie_draw
import model.data_loader as data_loader
import model.net as nn
import model.net_2 as nn2
import torch
import utils
from pic_draw.ie_digital import digit

if __name__=='__main__':
    # root path to
    acc_loss_dis_path = 'weights_params/IE demo val loss 60 epoch_random/'
    ie_path = 'epoch_params/60epoch/'
    print('acc')
    acc_curve, acc_label = accpic(acc_loss_dis_path)
    print('loss')
    loss_curve, loss_label =losspic(acc_loss_dis_path)

    json_path = 'experiments/base_cnn_2layer/params.json'
    params = utils.Params(json_path)
    params.cuda = torch.cuda.is_available()
    model = nn2.Net(params)
    ie_curve, ie_lsbel = ie_draw.ie_picnum(ie_path, model, )
    ie_curve, ie_label = digit()

    fig, (axes1, axes2, axes4) = plt.subplots(1, 3, figsize=(14, 5))

    plt.subplots_adjust(left=0.06, right=0.97, top=0.98, bottom=0.01, wspace=0.55, hspace=None)

    x_acc = []
    x_loss = []
    x_dis = []
    x_ie = []
    for i in range(len(acc_curve[0])):
        x_acc.append(i)
    for i in range(len(loss_curve[0])):
        x_loss.append(i)
    for i in range(len(ie_curve[0])):
        x_ie.append(i)

    for idx, acc in enumerate(acc_curve):
        axes1.plot(x_acc, acc, label=acc_label[idx])
    axes1.set_title('accuracy', fontsize = 17)
    axes1.set_xlabel('epoch', fontsize = 15)
    axes1.set_ylabel('accuracy', fontsize = 15)
    axes1.set_ylim([0.05, 0.75])
    axes1.set_aspect(aspect = 85.7, adjustable ='box')
    axes1.legend()

    for idx, loss in enumerate(loss_curve):
        axes2.plot(x_loss, loss, label=loss_label[idx])
    axes2.set_title('loss', fontsize = 17)
    axes2.set_xlabel('epoch', fontsize = 15)
    axes2.set_ylabel('loss', fontsize = 15)
    axes2.set_ylim([0.8, 2.5])
    axes2.set_aspect(aspect=35.2, adjustable='box')
    axes2.legend()

    for idx, ie in enumerate(ie_curve):
        axes4.plot(x_ie, ie, label=ie_label[idx])
    axes4.set_title('task relate information', fontsize = 17)
    axes4.set_xlabel('epoch', fontsize = 15)
    axes4.set_ylabel('task relate information', fontsize = 15)
    axes4.set_ylim([0, 1.75])
    axes4.set_aspect(aspect=34.28, adjustable='box')
    axes4.legend()

    plt.legend()
    plt.savefig('pic/pic2.pdf')
    plt.show()